#!/usr/bin/env python3
"""
WASM Unity Build Generator

Generates unity build files for faster WASM compilation by combining multiple
source files into fewer compilation units.

Unity builds work by creating wrapper .cpp files that #include other .cpp files.
This reduces:
  - Compiler invocations (250 files → ~6-10 unity files)
  - Duplicate header parsing overhead
  - Linking overhead (fewer object files)

Usage:
    python ci/wasm_unity_generator.py --chunks 10 --output build/wasm/unity
    python ci/wasm_unity_generator.py --list  # Just list the unity groups

Architecture:
    Input: List of source .cpp files from src/
    Output: Unity wrapper files (unity0.cpp, unity1.cpp, ...)

    Each unity file contains #include directives for a chunk of sources:
        #include </absolute/path/to/src/file1.cpp>
        #include </absolute/path/to/src/file2.cpp>
        ...

Performance:
    - Typical: 250 sources → 10 chunks = 25 sources per unity file
    - Compilation time: 60-80% faster for full rebuilds
    - Incremental builds: Slower (changing 1 file recompiles entire chunk)

    Recommendation: Use unity builds for CI/release, not daily development
"""

import argparse
import hashlib
import json
import sys
from pathlib import Path


# Project root
PROJECT_ROOT = Path(__file__).parent.parent


def get_source_files() -> list[Path]:
    """Get all .cpp source files from src/ directory."""
    src_dir = PROJECT_ROOT / "src"
    sources = []
    for cpp_file in src_dir.rglob("*.cpp"):
        # Exclude entry_point.cpp (compiled separately for sketches)
        if cpp_file.name == "entry_point.cpp":
            continue
        # Exclude test files
        if "test" in cpp_file.parts:
            continue
        sources.append(cpp_file.resolve())

    return sorted(sources)  # Sort for deterministic ordering


def chunk_sources(sources: list[Path], num_chunks: int) -> list[list[Path]]:
    """
    Split sources into roughly equal chunks.

    Args:
        sources: List of source file paths
        num_chunks: Number of unity chunks to create

    Returns:
        List of lists, where each sublist is a chunk of sources
    """
    if num_chunks <= 0:
        raise ValueError(f"num_chunks must be > 0, got {num_chunks}")

    if num_chunks > len(sources):
        num_chunks = len(sources)

    chunk_size = (len(sources) + num_chunks - 1) // num_chunks  # Ceiling division

    chunks = []
    for i in range(0, len(sources), chunk_size):
        chunk = sources[i : i + chunk_size]
        if chunk:  # Only add non-empty chunks
            chunks.append(chunk)

    return chunks


def generate_unity_file(
    chunk_index: int,
    sources: list[Path],
    output_path: Path,
) -> None:
    """
    Generate a single unity build file.

    Args:
        chunk_index: Index of this chunk (for naming)
        sources: List of source files to include in this unity file
        output_path: Path to write the unity file
    """
    lines = []
    lines.append(f"// Unity build file {chunk_index}")
    lines.append(f"// Generated by ci/wasm_unity_generator.py")
    lines.append(
        f"// Combines {len(sources)} source files into single compilation unit"
    )
    lines.append("")

    for source in sources:
        # Use absolute paths (Emscripten em++ handles them correctly)
        # Use angle brackets like Meson does
        lines.append(f"#include<{source.as_posix()}>")

    lines.append("")

    content = "\n".join(lines)

    # Only write if content changed (preserve timestamps for incremental builds)
    if output_path.exists():
        existing_content = output_path.read_text(encoding="utf-8")
        if existing_content == content:
            return  # No change, skip write

    output_path.write_text(content, encoding="utf-8")


def compute_unity_metadata_hash(num_chunks: int, sources: list[Path]) -> str:
    """Compute hash of unity configuration for change detection."""
    config_str = f"{num_chunks}:{len(sources)}:" + ":".join(
        str(s) for s in sorted(sources)
    )
    return hashlib.sha256(config_str.encode()).hexdigest()


def generate_unity_builds(
    num_chunks: int,
    output_dir: Path,
    verbose: bool = False,
) -> tuple[list[Path], dict[str, str | int]]:
    """
    Generate all unity build files.

    Args:
        num_chunks: Number of unity chunks to create
        output_dir: Directory to write unity files
        verbose: Enable verbose output

    Returns:
        Tuple of (list of unity file paths, metadata dict)
    """
    # Get all source files
    sources = get_source_files()

    if verbose:
        print(f"Found {len(sources)} source files")

    # Split into chunks
    chunks = chunk_sources(sources, num_chunks)

    if verbose:
        print(f"Split into {len(chunks)} chunks:")
        for i, chunk in enumerate(chunks):
            print(f"  Chunk {i}: {len(chunk)} files")

    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)

    # Generate unity files
    unity_files = []
    for i, chunk in enumerate(chunks):
        unity_path = output_dir / f"unity{i}.cpp"
        generate_unity_file(i, chunk, unity_path)
        unity_files.append(unity_path)

        if verbose:
            print(f"Generated: {unity_path.name} ({len(chunk)} sources)")

    # Compute metadata hash
    metadata = {
        "num_chunks": len(chunks),
        "num_sources": len(sources),
        "config_hash": compute_unity_metadata_hash(num_chunks, sources),
    }

    return unity_files, metadata


def list_unity_groups(num_chunks: int) -> None:
    """List what would be in each unity group without generating files."""
    sources = get_source_files()
    chunks = chunk_sources(sources, num_chunks)

    print(f"Unity build configuration:")
    print(f"  Total sources: {len(sources)}")
    print(f"  Unity chunks: {len(chunks)}")
    print()

    for i, chunk in enumerate(chunks):
        print(f"Unity chunk {i}: ({len(chunk)} files)")
        for source in chunk:
            rel_path = source.relative_to(PROJECT_ROOT)
            print(f"  - {rel_path}")
        print()


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Generate unity build files for WASM compilation"
    )
    parser.add_argument(
        "--chunks",
        type=int,
        default=10,
        help="Number of unity chunks to create (default: 10)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=PROJECT_ROOT / "build" / "wasm" / "unity",
        help="Output directory for unity files (default: build/wasm/unity)",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List unity groups without generating files",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Verbose output",
    )

    args = parser.parse_args()

    try:
        if args.list:
            list_unity_groups(args.chunks)
            return 0

        unity_files, metadata = generate_unity_builds(
            args.chunks,
            args.output,
            args.verbose,
        )

        # Write metadata
        metadata_path = args.output / "unity_metadata.json"
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=2)

        print(f"✓ Generated {len(unity_files)} unity build files")
        print(f"  Output: {args.output}")
        print(f"  Chunks: {metadata['num_chunks']}")
        print(f"  Sources: {metadata['num_sources']}")
        print(f"  Metadata: {metadata_path}")

        return 0

    except KeyboardInterrupt:
        print("\n✗ Interrupted by user")
        return 130
    except Exception as e:
        print(f"✗ Error: {e}", file=sys.stderr)
        import traceback

        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())
